{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving, TypeSynonymInstances, MultiParamTypeClasses #-}
-----------------------------------------------------------------------------
-- |
-- Module      : LLVM.Parser.Record.Parser
-- Copyright   : Mikhail Belyaev
-- License     : GPLv3 (see LICENSE)
--
--
-- TODO: description
-----------------------------------------------------------------------------
module LLVM.Parser.Record.Parser where 

import Control.Monad
import Control.Monad.State
import Control.Monad.Error
import Control.Monad.Writer

import LLVM.Parser.Bits.BitGet
import LLVM.Parser.Bits.Primitives
import LLVM.Parser.Bits.Struct
import LLVM.Parser.Record.Struct
import LLVM.Parser.Record.Record
import LLVM.Parser.Record.BlockInfo
import LLVM.Util

import qualified Data.ByteString as B
import Data.Map (Map,insert,fromList)
import qualified Data.Map as M
import Data.List (unfoldr) -- unfoldr
import Data.Maybe (fromJust,catMaybes)

import Data.PairLike
import Data.EitherLike
import Data.TreeLike

-- environment part of the state
data Environment = E {
                         records :: StructBindings, -- structure of records
                         curId :: Integer, -- current abbrev
                         curIdSize :: Int -- current id size
                     } deriving Show
                     
type EnvBindings = Map Int BlockInfo

defaultGlobal :: EnvBindings
defaultGlobal = fromList $ [ (0, blockInfoInfo) ]
-- parser state itself
data S = S {
                 position    :: Integer, -- number of records from the start
                 environment :: Environment,  -- parser environment
                 global :: EnvBindings -- global parser environment
               } deriving Show

-- increment position without affecting everything else
incS :: S -> S
incS input = let cur = position input in
             input{position = cur+1}

-- struct bindings of state S
bindsOf :: S -> StructBindings
bindsOf s = records $ environment s

-- current id size of state S
curIdSizeOf :: S -> Int
curIdSizeOf s = curIdSize $ environment s

-- add a binding to state (auto-assigning it an id and incrementing existing one)
addBind :: S -> [StructElem] -> S
addBind state bind = let env = environment state
                         nid = curId $ env
                         newBinds = insert nid bind $ bindsOf state
                     in state{environment = env{records = newBinds, curId = nid+1}}

-- add a binding to state (auto-assigning it an id and incrementing existing one)
addAbbr :: S -> AbbreviationDef -> S
addAbbr state (AbbreviationDef x) = addBind state x

addAbbrs :: [AbbreviationDef] -> S -> S
addAbbrs (a0:rest) state = addAbbrs rest (addAbbr state a0) 
addAbbrs [] state = state

setAbbrs :: [AbbreviationDef] -> S -> S
setAbbrs sum s = addAbbrs sum cleans
                 where env = environment s
                       cleans = s { environment = env{ records = defaultBinds, curId = 4 } }
                       
setIDSize :: Int -> S -> S
setIDSize alen s = s { environment = env{ curIdSize = fromIntegral alen } }
                   where env = environment s

-- the starting state for the parser
startingState :: S
-- default position = 0
-- default first id = 4
-- default starting id size = 2
startingState = S 0 (E defaultBinds 4 2) defaultGlobal

-- Parser monad, ErrBitGet + parser state
type Parser = StateT S ErrBitGet
-- Error-handling version of Parser monad
type ErrParser = ErrorOf' Parser

newtype ErrorOf' m a = EOF'{ unEOF':: ErrorOf m a} 
        deriving(Functor,Monad,MonadTrans,MonadPlus,MonadFix)

instance MonadState S ErrParser where
        get = lift get
        put = lift . put
        
instance MonadError String ErrParser where
        catchError wher how = EOF' $ catchError (unEOF' wher) (unEOF' . how)
        throwError mes = do
                                rec <- gets position
                                EOF' $ throwError $ show rec ++ ": " ++ mes

-- extract the error (if any) and lift the monad down the stack
-- one floor
unliftError' :: ErrParser a -> Parser (MayError a)
unliftError' p = unliftError $ unEOF' p

-- Is EOF reached?
isEof :: ErrParser Bool
isEof = lift . lift $ lift isEmpty

-- lift that gathers all errors in the topmost error monad
liftErr :: ErrBitGet a -> ErrParser a
liftErr = throwTo $ lift . lift

-- eval the parser on input, returns error if any or result
evalParser :: B.ByteString -> ErrParser a -> MayError a
evalParser input p = let parser = unliftError' p
                         errbitget = evalStateT parser startingState
                         bitget = unliftError errbitget
                     in
                     joinMe $ runBitGet input bitget

-- read a single record from the parser using inner struct bindings
readRecord :: ErrParser RawRecord
readRecord =  let 
                  getRecord'' x = liftErr $ getRecord (curIdSizeOf x) (bindsOf x) 
              in
              do
                  s <- get
                  rec <- getRecord'' s
                  put $ incS s
                  return rec

-- build a block from header and a list of records
toBlock :: BlockHeader -> [RawRecord] -> RawBlock
toBlock head body = RawBlock head $ unfoldr sliceBlockElem body

-- Partition the records by the next closing element on this level
-- the closing element is not included
-- ex.: [ent,clos,ent,clos,clos,ent,ent,ent]->[ent,clos,ent,clos][clos,ent,ent,ent]
cutFormed :: [RawRecord] -> Maybe ([RawRecord], [RawRecord])
cutFormed list = 
        let count ix (h:tail) ups downs 
                                     | downs > ups = Just (ix-1)
                                     | isBlockEntry h = count (ix+1) tail (ups+1) downs
                                     | isBlockClosing h = count (ix+1) tail ups (downs+1)
                                     | otherwise = count (ix+1) tail ups downs
            count ix [] ups downs 
                                     | downs > ups = Just (ix-1) 
                                     | otherwise = Nothing
        in case (count 0 list 0 0) of
                                Nothing -> Nothing
                                Just n -> Just $ splitAt n list

-- can be unfolded to proceed [RawRecord] -> [RawBlockElem] transform
-- the way it is, number of records and elems is (and must be) different
sliceBlockElem :: [RawRecord] -> Maybe (RawBlockElem,[RawRecord])
sliceBlockElem [] = Nothing
sliceBlockElem (r:rs) 
        | (not $ isBlockEntry r) && (not $ isBlockClosing r) = Just (Data r,rs)
        | isBlockEntry r = let check = cutFormed rs
                               head = fromRawRecord r -- gotta be just
                           in
                           case (check, head) of 
                                        (Nothing,Nothing) -> Nothing
                                        (Nothing,_) -> Nothing
                                        (_,Nothing) -> Nothing
                                        (Just (body,end:rest),Just chead)-> 
                                                Just (Sub $ toBlock chead body, rest)
        | otherwise = Nothing

-- take a record and, if it's an abbrevation def, consume it, fullfilling inner bindings
-- and return nothing
-- if it's not, return it as is
consumeAbbreviation :: RawRecord -> ErrParser (Maybe RawRecord)
consumeAbbreviation r | isAbbrDefinition r = do
                                                s <- get
                                                put $ s `addBind` (structure . fromJust $ fromRawRecord r)
                                                return Nothing
                      | otherwise = return $ Just r

getGlobalBinds :: Int -> ErrParser [AbbreviationDef]
getGlobalBinds id = let binfo glob = id `M.lookup` glob
                        binds Nothing = []
                        binds (Just b) = sbinds b
                    in gets (binds . binfo . global)

getBlockContext :: ErrParser Environment
getBlockContext = gets environment

putBlockContext :: Environment -> ErrParser ()
putBlockContext e = modify (\s -> s{environment = e})

-- put the idSize info from the header to the parser inner values
assumeBlockHead :: BlockHeader -> ErrParser ()
assumeBlockHead bh = do
                        s <- get
                        gbds <- getGlobalBinds $ fromInteger bid
                        modify $ setAbbrs gbds --`debug` (show bid ++ ":" ++ show gbds)
                        modify $ setIDSize alen
                        return ()
                     where 
                      alen = fromIntegral $ abbrLen bh
                      bid = blockId bh

assumeBlockInfos :: BlockInfoBlock -> ErrParser ()
assumeBlockInfos (BlockInfoBlock what) = 
                      do 
                         s <- get
                         glob <- gets global
                         put $ s{global = glob `M.union` what} --`debug` (show what)

---- raw block, an info header and a list of elems
--data RawBlock = RawBlock { header :: BlockHeader, contents :: [RawBlockElem] } deriving Show
---- block elem, a sub elem or a record
--data RawBlockElem = Data RawRecord | Sub RawBlock deriving Show




